﻿
#include "alpc.h"
#include <Windows.h>

PFN_NtAlpcCreatePort NtAlpcCreatePort = NULL;
PFN_NtAlpcConnectPort NtAlpcConnectPort = NULL;
PFN_NtAlpcConnectPortEx NtAlpcConnectPortEx = NULL;
PFN_NtAlpcAcceptConnectPort NtAlpcAcceptConnectPort = NULL;
PFN_NtAlpcDisconnectPort NtAlpcDisconnectPort = NULL;
PFN_NtAlpcSetInformation NtAlpcSetInformation = NULL;
PFN_NtAlpcQueryInformation NtAlpcQueryInformation = NULL;
PFN_NtAlpcSendWaitReceivePort NtAlpcSendWaitReceivePort = NULL;
PFN_NtAlpcCreateSecurityContext NtAlpcCreateSecurityContext = NULL;
PFN_NtAlpcCancelMessage NtAlpcCancelMessage = NULL;
PFN_NtAlpcCreatePortSection NtAlpcCreatePortSection = NULL;

BOOL InitAlpcEnv()
{
    HMODULE hNtDll = GetModuleHandleW(L"ntdll.dll");
    if (NULL == hNtDll || INVALID_HANDLE_VALUE == hNtDll)
    {
        return NULL;
    }

#define X(Name) ((*(FARPROC *)&Name = GetProcAddress(hNtDll, #Name)) == NULL)
    X(NtAlpcCreatePort);
    X(NtAlpcConnectPort);
    X(NtAlpcConnectPortEx);
    X(NtAlpcAcceptConnectPort);
    X(NtAlpcDisconnectPort);
    X(NtAlpcSetInformation);
    X(NtAlpcQueryInformation);
    X(NtAlpcSendWaitReceivePort);
    X(NtAlpcCreateSecurityContext);
    X(NtAlpcCancelMessage);
    X(NtAlpcCreatePortSection);
    return !!hNtDll;
}

void UninitAlpcEnv()
{

}


VOID NTAPI AlpclInitUnicodeString(
    PUNICODE_STRING DestinationString,
    PCWSTR SourceString
)
{
    // 设置 UNICODE_STRING 的字段
    DestinationString->Buffer = (PWSTR)SourceString;
    DestinationString->Length = (USHORT)(wcslen(SourceString) * sizeof(WCHAR));  // 计算实际长度（不包括 NULL 字符）
    DestinationString->MaximumLength = (USHORT)((wcslen(SourceString) + 1) * sizeof(WCHAR)); // 计算最大长度（包括 NULL 字符）
}

HANDLE AlpcCreatePort(LPCWSTR strName)
{
    NTSTATUS Status = 0;
    HANDLE hPort = NULL;
    UNICODE_STRING strPortName = { 0 };
    OBJECT_ATTRIBUTES ObjectAttributes = { 0 };
    ALPCPORT_ATTRIBUTES PortAttributes = { 0 };


    // 初始化 ALPC 端口属性
    PortAttributes.MaxMessageLength = ALPC_MAX_BUFLEN;
    PortAttributes.MemoryBandwidth = 512;
    PortAttributes.MaxPoolUsage = 0xffffffff;
    PortAttributes.MaxSectionSize = 0xffffffff;
    PortAttributes.MaxViewSize = 0xffffffff;
    PortAttributes.MaxTotalSectionSize = 0xffffffff;
    PortAttributes.Flags = ALPC_PORTFLG_ALLOW_DUP_OBJECT | ALPC_PORTFLG_ALLOWIMPERSONATION 
        | ALPC_PORTFLG_LRPC_WAKE_POLICY1 | ALPC_PORTFLG_LRPC_WAKE_POLICY2 | ALPC_PORTFLG_LRPC_WAKE_POLICY3;

    // 初始化 ALPC 安全设置 (可选)
    SECURITY_QUALITY_OF_SERVICE SecurityQos = { 0 };
    SecurityQos.ImpersonationLevel = SecurityIdentification;
    SecurityQos.ContextTrackingMode = SECURITY_STATIC_TRACKING;
    SecurityQos.EffectiveOnly = 0;
    SecurityQos.Length = sizeof(SecurityQos);
    //PortAttributes.SecurityQos = SecurityQos;

    if (NULL == strName || L'\0' == strName[0])
    {
        return NULL;
    }

    // 初始化端口名称
    AlpclInitUnicodeString(&strPortName, strName);

    // 设置对象属性
    InitializeObjectAttributes(
        &ObjectAttributes,
        &strPortName,
        0,
        NULL,
        NULL
    );

    Status = NtAlpcCreatePort(&hPort, &ObjectAttributes, &PortAttributes);

    if (!ALPCNT_SUCCESS(Status)) {
        return NULL;
    }

    return hPort;
}


HANDLE AlpcConnectPort(LPCWSTR strName)
{
    // 连接到服务端
    NTSTATUS Status = 0;
    HANDLE hPort = NULL;
    // 端口名称（与服务端相同）
    UNICODE_STRING strPortName = { 0 };
    SECURITY_QUALITY_OF_SERVICE SecurityQos = { 0 };
    ALPCPORT_ATTRIBUTES PortAttributes = { 0 };

    ALPCPORT_MESSAGE ConnectMessage = { 0 };
    ConnectMessage.u1.s1.DataLength = 0;  // 无数据
    ConnectMessage.u1.s1.TotalLength = sizeof(ALPCPORT_MESSAGE);
    ConnectMessage.MessageId = 0;
    DWORD dwDataSize = sizeof(ALPCPORT_MESSAGE);

    // QOS
    SecurityQos.ImpersonationLevel = SecurityImpersonation; // SecurityIdentification; // ;
    SecurityQos.ContextTrackingMode = SECURITY_STATIC_TRACKING;
    SecurityQos.EffectiveOnly = 0;
    SecurityQos.Length = sizeof(SecurityQos);
    // ALPC Port Attributs
    PortAttributes.Flags = ALPC_PORTFLG_ALLOW_DUP_OBJECT | ALPC_PORTFLG_ALLOWIMPERSONATION | ALPC_PORTFLG_LRPC_WAKE_POLICY1 | ALPC_PORTFLG_LRPC_WAKE_POLICY2 | ALPC_PORTFLG_LRPC_WAKE_POLICY3; // | ALPC_PORTFLG_DIRECT_MESSAGE;
    PortAttributes.MaxMessageLength = ALPC_MAX_BUFLEN; // technically the hard limit for this is 65535, if no constrains you can use AlpcMaxAllowedMessageLength() to set this limit
    PortAttributes.MemoryBandwidth = 512;
    PortAttributes.MaxSectionSize = 0xffffffff; // 20000; 
    PortAttributes.MaxViewSize = 0xffffffff; // 20000; // sizeof(PORT_VIEW); 
    PortAttributes.MaxTotalSectionSize = 0xffffffff;// 20000;
    PortAttributes.DupObjectTypes = 0xffffffff;
    PortAttributes.MaxPoolUsage = 0xffffffff; // 0x4000;
    PortAttributes.SecurityQos = SecurityQos;

    if (NULL == strName || L'\0' == strName[0])
    {
        return NULL;
    }

    // 初始化端口名称
    AlpclInitUnicodeString(&strPortName, strName);

    Status = NtAlpcConnectPort(
        &hPort,                    // 返回的端口句柄
        &strPortName,           // 服务端端口名称
        NULL,                      // 客户端对象属性
        &PortAttributes,                      // 端口属性
        ALPC_SYNC_CONNECTION,      // 连接标志
        NULL,                      // 安全描述符
        &ConnectMessage, // 连接消息
        NULL,// 连接消息大小      
        NULL,                      // 输出消息属性
        NULL,                      // 输入消息属性
        NULL                   // 超时时间
    );

    if (!ALPCNT_SUCCESS(Status)) {
        return NULL;
    }

    return hPort;
}

HANDLE AlpcAcceptPort(HANDLE hServer, ULONG uMessageId)
{
    // 接受连接请求
    NTSTATUS Status = 0;
    HANDLE hClient = NULL;
    ALPCPORT_MESSAGE RecvData = { 0 };
    RecvData.u1.s1.DataLength = 0;  // 无数据
    RecvData.u1.s1.TotalLength = sizeof(ALPCPORT_MESSAGE);
    RecvData.MessageId = uMessageId;

    Status = NtAlpcAcceptConnectPort(
        &hClient,
        hServer,
        ALPC_SYNC_CONNECTION,
        NULL,
        //&PortAttributes,
        NULL,
        NULL,
        &RecvData,
        //&RecvDataAttr,
        NULL,
        TRUE
    );

    if (!ALPCNT_SUCCESS(Status)) {
        return NULL;
    }

    return hClient;
}

VOID AlpcClose(HANDLE hPort)
{
    NTSTATUS Status = 0;
    Status = NtAlpcDisconnectPort(hPort, 0);
    CloseHandle(hPort);
}

BOOL AlpcSendWaitReceive(HANDLE hPort, PALPCPORT_MESSAGE pSendData,
    PALPCPORT_MESSAGE pRecvData, DWORD* pdwRecvSize, DWORD dwTimeOut)
{
    NTSTATUS Status = 0;
    ALPCPORT_ATTRIBUTES PortAttributes = { 0 };
    LARGE_INTEGER Timeout = { 0 };
    Timeout.QuadPart = -10000 * dwTimeOut;  // 1毫秒超时

    Status = NtAlpcSendWaitReceivePort(
        hPort,                         // 端口句柄
        0,
        pSendData,                       // 发送消息
        NULL,                               // 发送消息的附加属性
        pRecvData,                    // 接收消息
        pdwRecvSize,                      // 接收缓冲区长度
        &PortAttributes,                               // 接收消息的附加属性
        -1 == dwTimeOut ? NULL : &Timeout                             // 超时设置
    );

    if (!ALPCNT_SUCCESS(Status)) {
        return FALSE;
    }

    return TRUE;
}

BOOL AlpcClientSendReceive(LPCWSTR strName, PALPCPORT_MESSAGE pSendData,
    PALPCPORT_MESSAGE pRecvData, DWORD* pdwRecvSize, DWORD dwTimeOut)
{
    BOOL bRet = FALSE;
    HANDLE hPort = AlpcConnectPort(strName);
    if (NULL == hPort)
    {
        return FALSE;
    }

    bRet = AlpcSendWaitReceive(hPort, pSendData, pRecvData, pdwRecvSize, dwTimeOut);

    AlpcClose(hPort);

    return bRet;
}